package cn.actoncode.boot.framework.zhipu.core.client;

import cn.actoncode.boot.framework.common.exception.ServiceException;
import cn.actoncode.boot.framework.zhipu.config.ZhipuProperties;
import cn.actoncode.boot.framework.zhipu.core.bean.*;
import cn.actoncode.boot.framework.zhipu.core.enums.ChatGLMRoleEnum;
import cn.actoncode.boot.framework.zhipu.core.utils.OkHttpUtils;
import cn.hutool.core.util.ArrayUtil;
import cn.hutool.core.util.IdUtil;
import com.alibaba.fastjson2.JSON;
import jakarta.annotation.PostConstruct;
import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.dromara.hutool.json.jwt.JWTUtil;

import java.io.IOException;
import java.io.InputStream;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

import static cn.actoncode.boot.framework.common.exception.enums.GlobalErrorCodeConstants.*;
import static cn.actoncode.boot.framework.common.exception.util.ServiceExceptionUtil.exception;

@Slf4j
@AllArgsConstructor
public class ChatGLMApiClient {

    public static final int [] TOKEN_ERROR_CODE = {401, 1003, 1004};
    public static final int [] SEC_CHECK_ERROR_CODE = {1301, 1505};
    public static final int [] QPS_ERROR_CODE = {1300, 1302, 1303, 1304, 1305, 429, 503};

    @Getter
    private final ZhipuProperties zhipuProperties;

    private volatile String token;

    public ChatGLMApiClient(ZhipuProperties zhipuProperties) {
        this.zhipuProperties = zhipuProperties;
    }

    @PostConstruct
    public void init() {
        // 刷新token
        refreshToken();
    }

    public void refreshToken() {
        Map<String, Object> headers = new HashMap<>(2);
        headers.put("alg", "HS256");
        headers.put("sign_type", "SIGN");

        Map<String, Object> payload = new HashMap<>(2);
        payload.put("api_key", zhipuProperties.getApiKeyId());
        payload.put("exp", System.currentTimeMillis() + zhipuProperties.getExp() * 1000);
        payload.put("timestamp", System.currentTimeMillis());
        token = JWTUtil.createToken(headers, payload, zhipuProperties.getApiKeySecret().getBytes());
    }

    public ChatGLMV4Response chat(ChatGLMV4Request chatGLMRequest, String api) {

        String requestStr = JSON.toJSONString(chatGLMRequest);
        log.info("ChatGlm4: {}", requestStr);
        Request request = new Request.Builder()
                .url(api)
                .post(RequestBody.create(requestStr, MediaType.parse("application/json; charset=utf-8")))
                .header("Authorization", this.token)
                .header("Content-type", "application/json; charset=utf-8")
                .build();

        Call call = OkHttpUtils.getInstance().newCall(request);
        call.timeout().timeout(180, TimeUnit.SECONDS);
        try {
            Response response = call.execute();
            String responseBody = response.body().string();
            if (response.code() == 200) {
                return JSON.parseObject(responseBody, ChatGLMV4Response.class);
            } else {
                ChatGLMStreamError chatGLMStreamError = JSON.parseObject(responseBody, ChatGLMStreamError.class);
                ChatGLMStreamError.Error error = chatGLMStreamError.getError();
                int code = error.getCode();
                if (ArrayUtil.contains(TOKEN_ERROR_CODE, code)){
                    refreshToken();
                    return chat(chatGLMRequest, api);
                } else {
                    throw getServiceException(responseBody, error.getMessage(), code);
                }
            }
        } catch (IOException e) {
            log.info("ChatGLM Error: {}", e.getMessage());
            throw exception(ZHIPU_AI_CHAT_ERROR);
        }
    }

    public InputStream streamChat(ChatGLMV4Request chatGLMV4Request, String url) {
        log.info("ChatGlm4: {}", JSON.toJSONString(chatGLMV4Request));
        Request request = new Request.Builder()
                .url(url)
                .post(RequestBody.create(JSON.toJSONString(chatGLMV4Request), MediaType.parse("application/json; charset=utf-8")))
                .header("Authorization", this.token)
                .header("Content-type", "text/event-stream")
                .build();

        Call call = OkHttpUtils.getInstance().newCall(request);
        call.timeout().timeout(180, TimeUnit.SECONDS);
        try {
            Response response = call.execute();
            if (response.code() == 200) {
                return response.body().byteStream();
            } else {
                String responseBody = response.body().string();
                ChatGLMStreamError chatGLMStreamError = JSON.parseObject(responseBody, ChatGLMStreamError.class);
                ChatGLMStreamError.Error error = chatGLMStreamError.getError();
                int code = error.getCode();
                if (ArrayUtil.contains(TOKEN_ERROR_CODE, code)){
                    refreshToken();
                    return streamChat(chatGLMV4Request, url);
                } else {
                    throw getServiceException(responseBody, error.getMessage(), code);
                }
            }
        } catch (IOException e) {
            log.info("ChatGLM Error: {}", e.getMessage());
            throw exception(ZHIPU_AI_CHAT_ERROR);
        }
    }

    private ServiceException getServiceException(String responseBody, String message, int code) {
        if (ArrayUtil.contains(SEC_CHECK_ERROR_CODE, code)) {
            return exception(NLP_SEC_CHECK_ERROR);
        } else if (ArrayUtil.contains(QPS_ERROR_CODE, code)){
            return exception(ZHIPU_AI_QPS_ERROR);
        } else {
            log.info("ChatGLM INFO: {}", responseBody);
            log.info("ChatGLM Error: {}", message);
            return exception(ZHIPU_AI_CHAT_ERROR);
        }
    }

    public ChatGLMRequest buildRequest(String prompt, List<ChatGLMPrompt> messages) {

        ChatGLMPrompt chatGLMPrompt = new ChatGLMPrompt();
        chatGLMPrompt.setRole(ChatGLMRoleEnum.user.name());
        chatGLMPrompt.setContent(prompt);
        messages.add(chatGLMPrompt);
        // 构建请求
        ChatGLMRequest chatGLMRequest = new ChatGLMRequest();
        chatGLMRequest.setRequestId(IdUtil.getSnowflakeNextIdStr());
        chatGLMRequest.setPrompt(messages);
        chatGLMRequest.setTemperature(0.95);
        chatGLMRequest.setTopP(0.7);
        return chatGLMRequest;
    }

    public ChatGLMV4Request buildRequestV4(String prompt, Boolean stream, List<ChatGLMV4Request.Message> messages) {

        ChatGLMV4Request.Message message = new ChatGLMV4Request.Message();
        message.setRole(ChatGLMRoleEnum.user.name());
        message.setContent(prompt);
        messages.add(message);
        // 构建请求
        ChatGLMV4Request chatGLMRequest = new ChatGLMV4Request();
        chatGLMRequest.setRequestId(IdUtil.getSnowflakeNextIdStr());
        chatGLMRequest.setStream(stream);
        chatGLMRequest.setMessages(messages);
        chatGLMRequest.setTemperature(0.95);
        chatGLMRequest.setTopP(0.7);
        return chatGLMRequest;
    }

}
